import torch


def calculate_acc(logits, labels, ignore_index=-100):
    """
    计算当前 batch 的预测正确的 token 数量 和 有效 token 的总数（忽略 ignore_index 位置）
    - 考虑 language modeling 的错位预测（input[t] -> predict[t+1]）
    - logits: [batch_size, seq_len, vocab_size]
    - labels: [batch_size, seq_len]
    """
    # 去掉 logits 最后一位（预测不到）
    logits = logits[:, :-1, :]  # shape: (batch, seq_len-1, vocab)

    # 去掉 labels 第一位（CLS），使其与 logits 对齐
    labels = labels[:, 1:]      # shape: (batch, seq_len-1)


    # 获取预测结果（token id）
    pred_ids = torch.argmax(logits, dim=-1)  # shape: (batch, seq_len-1)

    # 只计算有效位置（非 ignore_index）
    mask = labels != ignore_index

    # 正确且有效的预测位置
    correct = (pred_ids == labels) & mask

    correct_num = correct.sum().item()
    total_num = mask.sum().item()

    return correct_num, total_num


